'''
Created on Sep 20, 2011

@author: oabalbin
'''

'''
Created on Nov 17, 2010

@author: oabalbin
'''
import sys
import os
import glob
from optparse import OptionParser

from exome.jobs.base import JOB_SUCCESS, JOB_ERROR
from exome.jobs.job_runner import qsub_cac, qsub_loc, run_local
from exome.jobs.config import ExomePipelineConfig, ExomeAnalysisConfig
from collections import defaultdict, deque

# Global Variables
NODE_MEM=45000.0
NODE_CORES=12
SINGLE_CORE=1
MEM_PER_CORE= int(float(NODE_MEM) / NODE_CORES)
# wt=walltime
WT_SHORT= "24:00:00"
WT_LONG= "100:00:00"#"60:00:00" #"100:00:00"


def check_create_dir(root_path, dir_name=None):
    if not os.path.isdir(root_path):
            os.mkdir( root_path )
    if dir_name is not None:
        subfolder=os.path.join(root_path,dir_name)
        if not os.path.isdir(subfolder):
                os.mkdir( subfolder )
    else:             
        subfolder=root_path

    return subfolder


def read_files_folder(folderpath,ext):
    ''' '''
    # Read files in folder
    myfiles=[]
    for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
        myfiles.append(infile)
        #print "current file is: " + infile
    return myfiles


def samtools_mpileup(ref_genome, list_bam_files, outfile, path_to_sam, target_regions=None,call_parameters=None):
    '''
    Returns the shell to execute a pile up using sam tools 
    param is a dictionary with the input param
    '''
    samtools_command=path_to_sam+'samtools'
    bcftools_command=path_to_sam+'bcftools/bcftools'
    outtem=open(outfile,'w')
    outtem.close()

    
    default_mapping_quality=str(20)
    defualt_mappling_quality=str(13)
    # extra parameters
    if call_parameters is not None:
        min_mapping_qual = call_parameters[0]
        min_base_qual = call_parameters[1]

    else:
        min_mapping_qual =  default_mapping_quality
        min_base_qual = defualt_mappling_quality

    '''
    -C50 # QUlity downgrading specifically for BWA
    -q # mapping quality
    -Q # Base Quality
    '''
    args0=['-q',min_mapping_qual,'-Q',min_base_qual]

    if target_regions is not None:
        args=[samtools_command,'mpileup']+ args0+['-f',ref_genome] + list_bam_files +['-l', target_regions, '>', outfile]
    else:
        args=[samtools_command,'mpileup']+ args0+['-f',ref_genome] + list_bam_files +['>', outfile]
        
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command

def bcftools_call_snps(bcf_file_name, outfile_name, path_to_sam):
    '''
    Call snps candidates using bcftools
    This tool converts BCF to VCF, indexes BCF for random access, concatenates (not merges) BCFs, 
    estimate site allele frequencies and calls SNP candidates.
    '''
    bcftools_command=path_to_sam+'bcftools/bcftools'
    args = [bcftools_command, 'view', '-cg', bcf_file_name, '>', outfile_name]
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def sam_bcftools_mpileup(ref_genome, list_bam_files, outfile, path_to_sam, target_regions=None, call_parameters=None):
    '''
    '''
    samtools_command=path_to_sam+'samtools'
    bcftools_command=path_to_sam+'bcftools/bcftools'
    outtem=open(outfile,'w')
    outtem.close()
    
    default_mapping_quality=str(20)
    defualt_mappling_quality=str(13)
    # extra parameters
    if call_parameters is not None:
        min_mapping_qual = call_parameters[0]
        min_base_qual = call_parameters[1]

    else:
        min_mapping_qual =  default_mapping_quality
        min_base_qual = defualt_mappling_quality

    '''
    -C50 # QUlity downgrading specifically for BWA
    -q # mapping quality
    -Q # Base Quality
    '''
    args0=['-q',min_mapping_qual,'-Q',min_base_qual]
    if target_regions is not None:
        args=[samtools_command,'mpileup']+ args0+['-uf',ref_genome]+ list_bam_files +['-l', target_regions, '|'] +\
        [bcftools_command, 'view', '-vcg', '-', '>', outfile]
    else:
        args=[samtools_command,'mpileup']+ args0+['-uf',ref_genome] + list_bam_files +['|']+\
        [bcftools_command, 'view', '-vcg', '-', '>', outfile]
        
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command

       
def snps_calling_pairedSamples(analysis, configrun, use_recal_files, jobrunfunc, prev_deps):
    
    ####    
    extra_mem, num_cores = configrun.gatk_use_mem, configrun.gatk_num_cores
    path_to_vcftools, path_to_picard, path_to_sam = configrun.vcftools_path, configrun.picard_path, \
                                                configrun.samtools_path
    target_exons = configrun.sam_target_exons
    print target_exons
    genomes = configrun.genomes['human']
    ref_genome = genomes.gatk_ref_genome
    my_email=configrun.email_addresses
    
    sam_call_parameters=configrun.sam_mpileup_params
    cohort_samples = defaultdict(deque)
    
    if use_recal_files:
        analysis.sam_tools_files(use_recal_files)
        for sp in analysis.samples:
            if sp.category=='benign':
                cohort_samples['benign'].append(sp.sorted_mmarkdup_bam)
            else:
                cohort_samples['tumor'].append(sp.sorted_mmarkdup_bam)
    else:
        analysis.sam_tools_files(use_recal_files)
        for sp in analysis.samples:
            if sp.category=='benign':
                cohort_samples['benign'].append(sp.sorted_quickmmdup_bam)
            else:
                cohort_samples['tumor'].append(sp.sorted_quickmmdup_bam)


    check_create_dir(analysis.sam_calls_dir)
    deps_list=[]
    
    
    for sptype, list_bam_samples in cohort_samples.iteritems():
        #
        jobn='sst'
        jobn=jobn+sptype
        
        if sptype == 'benign':
            bcf_file_name=analysis.benign_bcf_mpileup_file
            vcf_file_name=analysis.benign_mpileup_file
        elif sptype == 'tumor':
            bcf_file_name=analysis.tumor_bcf_mpileup_file
            vcf_file_name=analysis.tumor_mpileup_file
        
        '''
        # Compute first the mpileup, then call the snvs using gatk and varscan.
        command = samtools_mpileup(ref_genome, list(list_bam_samples), bcf_file_name, path_to_sam, target_exons,sam_call_parameters)
        # Create the bcf mpileup file
        jobidbcf = jobrunfunc(jobn+'_bcf', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                           deps=prev_deps, stdout=None, email_addresses=my_email)
        
        # Create the vcf file of candidate snps
        command = bcftools_call_snps(bcf_file_name, vcf_file_name, path_to_sam)
        jobidbcf2 = jobrunfunc(jobn+'_vcf', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                           deps=jobidbcf, stdout=None, email_addresses=my_email)

        '''
        

        ########
        command = sam_bcftools_mpileup(ref_genome, list(list_bam_samples), vcf_file_name, path_to_sam, target_exons,sam_call_parameters)
        # Create the vcf file of candidate snps
        jobidbcf2 = jobrunfunc(jobn+'_vcf', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                           deps=prev_deps, stdout=None, email_addresses=my_email)
        deps_list.extend([jobidbcf2])
    
    return deps_list 

        

def snps_calling_multipleSamples(analysis, config, num_processors, jobrunfunc, prev_deps):
    print 'This is function is not implemented yet. It is work in progress'
    

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-r", "--config_file", dest="config_file",
                            help="file with run configuration")
    optionparser.add_option("-a", "--analysis_file", dest="analysis_file",
                            help="file with experiment configuration") 
    optionparser.add_option("--paired_samples", dest="paired_samples", action="store_true", default=False,
                            help="paired samples snv calling") 
    optionparser.add_option("--multi_samples", dest="multi_samples", action="store_true", default=False,
                            help="multi-sample snv calling")
    optionparser.add_option("--use_recal_files", dest="use_recal_files", action="store_true", default=False,
                            help="use recal files to call the snvs variants")

      
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    optionparser.add_option("--local_cluster", dest="local_cluster", action="store_true", default=False)
    optionparser.add_option("-p", "--processes", type=int, dest="num_processors", default=1)

    (options, args) = optionparser.parse_args()    

    config = ExomePipelineConfig()
    config.from_xml(options.config_file)
    analysis = ExomeAnalysisConfig()
    analysis.from_xml(options.analysis_file, config.output_dir)
    #Default when called from the command line
    
    depends=[]
            
    if not (options.local ^ options.cluster ^ options.local_cluster):
        optionparser.error("Must set either --local or --cluster to run job")
    if options.local:
        jobrunfunc = run_local
    elif options.cluster:
        jobrunfunc = qsub_cac
    elif options.local_cluster:
        jobrunfunc = qsub_loc

    
    if options.multi_samples:
        # Work in progress. 12-14-10
        snps_calling_multipleSamples(analysis, config, options.use_recal_files, jobrunfunc,depends)
    elif options.paired_samples:        
        snps_calling_pairedSamples(analysis, config, options.use_recal_files, jobrunfunc, depends)

    